# -*- coding: utf-8 -*-
"""
🚀 高性能异步 TCP 压测脚本（支持 6 种场景）
Author: chen xi
Date: 2025/9/11
支持模式：
1. request-response       -> 发送 → 接收 → 断开
2. echo-loop              -> 发送 → 接收 → 再发送（循环）
3. streaming-receive      -> 持续接收数据，直到服务端断开（长连接）
4. streaming-send         -> 持续发送数据
5. connect-and-close      -> 连接后立即断开
6. heartbeat-long-connection -> 长连接 + 心跳保活（新增）
"""

import asyncio
import time
import statistics
import sys
from typing import Tuple, Callable

# =================== 配置项 ===================
SERVER_IP = '47.107.56.10'
SERVER_PORT = 50002

# 压测模式（选择其一）：
#   'request-response'     : 发送 → 接收 → 断开
#   'echo-loop'            : 发送 → 接收 → 再发送（循环 3 次）
#   'streaming-receive'    : 连接后持续接收，等待服务端超时断开
#   'streaming-send'       : 连接后持续发送 10 秒
#   'connect-and-close'    : 连接后立即断开
#   'heartbeat-long-connection': 长连接 + 定期发送心跳
TEST_MODE = 'echo-loop'

# 数据与并发
MESSAGE = b'0123456789'  # 发送的数据
NUM_WORKERS = 10000  # 并发协程数（每个协程一个长连接）
DURATION = 3600  # 测试总时长（秒），建议 ≥ 服务端超时时间
STAT_INTERVAL = 10  # 统计输出间隔（秒）
TIMEOUT = 5  # 网络操作超时（秒）
RECV_SIZE = 1024  # 接收缓冲区大小
ECHO_LOOP_COUNT = 3  # echo-loop 模式下的循环次数
STREAMING_DURATION = 10  # streaming-send 模式的发送时长

# --- 心跳配置（仅 heartbeat-long-connection 模式使用）---
HEARTBEAT_INTERVAL = 3600  # 心跳间隔（秒）
HEARTBEAT_MESSAGE = b'PING\n'  # 心跳包内容，根据服务端协议调整
# 示例：b'{"cmd":"ping"}\n' 或 b'\x01\x00\x00\xFF'

# =============================================

# 全局统计
total_success = 0
total_fail = 0
total_latencies = []
total_bytes_received = 0

# 当前窗口统计
current_success = 0
current_fail = 0
current_latencies = []
current_bytes = 0

# 锁
lock = asyncio.Lock()

# 控制变量
running = True


# =================== 各模式的行为函数 ===================

async def behavior_request_response(reader, writer) -> Tuple[bool, float, int]:
	"""场景1: 发送 → 接收 → 断开"""
	start_time = time.time()
	try:
		if MESSAGE:
			writer.write(MESSAGE)
			await writer.drain()
		
		response = await asyncio.wait_for(reader.read(RECV_SIZE), timeout=TIMEOUT)
		return True, (time.time() - start_time) * 1000, len(response)
	except Exception as e:
		print(f"❌ [request-response] {type(e).__name__}: {e}", file=sys.stderr)
		return False, 0, 0


async def behavior_echo_loop(reader, writer) -> Tuple[bool, float, int]:
	"""场景2: 发送 → 接收 → 再发送（循环 N 次）"""
	start_time = time.time()
	total_bytes = 0
	try:
		for i in range(ECHO_LOOP_COUNT):
			if MESSAGE:
				writer.write(MESSAGE)
				await writer.drain()
			
			response = await asyncio.wait_for(reader.read(RECV_SIZE), timeout=TIMEOUT)
			total_bytes += len(response)
		
		return True, (time.time() - start_time) * 1000, total_bytes
	except Exception as e:
		print(f"❌ [echo-loop] {type(e).__name__}: {e}", file=sys.stderr)
		return False, 0, 0


async def behavior_streaming_receive(reader, writer) -> Tuple[bool, float, int]:
    """场景3: 持续接收数据，断开后自动重连"""
    total_bytes = 0
    start_time = time.time()  # 记录首次连接开始时间

    while running:
        try:
            # 持续接收
            while running:
                data = await asyncio.wait_for(reader.read(RECV_SIZE), timeout=1.0)
                if not data:
                    break  # 服务端正常关闭
                total_bytes += len(data)

            # 被动断开，准备重连
            print(f"🔁 [streaming] 连接已断开，正在重连...")
            await asyncio.sleep(0.1)  # 避免忙等

            # 重连
            reader, writer = await asyncio.wait_for(
                asyncio.open_connection(SERVER_IP, SERVER_PORT),
                timeout=TIMEOUT
            )
            print(f"✅ [streaming] 重连成功")

        except (ConnectionResetError, ConnectionAbortedError, BrokenPipeError):
            # 连接被对端重置或中断
            print(f"🔁 [streaming] 连接中断，正在重连...")
            await asyncio.sleep(0.1)
            try:
                reader, writer = await asyncio.wait_for(
                    asyncio.open_connection(SERVER_IP, SERVER_PORT),
                    timeout=TIMEOUT
                )
                print(f"✅ [streaming] 重连成功")
            except Exception as e:
                print(f"❌ [streaming] 重连失败: {e}")
                return False, (time.time() - start_time) * 1000, total_bytes

        except asyncio.TimeoutError:
            print(f"⏰ [streaming] 接收超时，可能网络异常")
            await asyncio.sleep(1)
            continue

        except Exception as e:
            print(f"❌ [streaming] 未知错误: {type(e).__name__}: {e}")
            return False, (time.time() - start_time) * 1000, total_bytes

    return True, (time.time() - start_time) * 1000, total_bytes


async def behavior_streaming_send(reader, writer) -> Tuple[bool, float, int]:
	"""场景4: 持续发送数据（固定时长）"""
	start_time = time.time()
	total_bytes = 0
	end_time = start_time + STREAMING_DURATION
	try:
		while time.time() < end_time and running:
			if MESSAGE:
				writer.write(MESSAGE)
				await writer.drain()
				total_bytes += len(MESSAGE)
			await asyncio.sleep(0.001)
		return True, (time.time() - start_time) * 1000, total_bytes
	except Exception as e:
		print(f"❌ [streaming-send] {type(e).__name__}: {e}", file=sys.stderr)
		return False, 0, 0


async def behavior_connect_and_close(reader, writer) -> Tuple[bool, float, int]:
	"""场景5: 连接后立即断开（不发不收）"""
	start_time = time.time()
	try:
		return True, (time.time() - start_time) * 1000, 0
	except Exception as e:
		print(f"❌ [connect-and-close] {type(e).__name__}: {e}", file=sys.stderr)
		return False, 0, 0


async def behavior_heartbeat_long_connection(reader, writer) -> Tuple[bool, float, int]:
    """场景6: 长连接 + 心跳保活 + 断线重连"""
    total_bytes = 0
    start_time = time.time()
    last_heartbeat = time.time()

    while running:
        try:
            # 主循环：接收 + 发送心跳
            while running:
                # 接收数据
                try:
                    data = await asyncio.wait_for(reader.read(RECV_SIZE), timeout=1.0)
                    if not data:
                        break
                    total_bytes += len(data)
                except asyncio.TimeoutError:
                    pass
                except (ConnectionResetError, ConnectionAbortedError):
                    raise  # 触发外层重连

                # 发送心跳
                if HEARTBEAT_MESSAGE and (time.time() - last_heartbeat) >= HEARTBEAT_INTERVAL:
                    writer.write(HEARTBEAT_MESSAGE)
                    await writer.drain()
                    total_bytes += len(HEARTBEAT_MESSAGE)
                    last_heartbeat = time.time()

            # 正常断开，准备重连
            print(f"🔁 [heartbeat] 服务端关闭连接，正在重连...")
            await asyncio.sleep(0.1)

            reader, writer = await asyncio.wait_for(
                asyncio.open_connection(SERVER_IP, SERVER_PORT),
                timeout=TIMEOUT
            )
            print(f"✅ [heartbeat] 重连成功")
            last_heartbeat = time.time()

        except (ConnectionResetError, ConnectionAbortedError, BrokenPipeError):
            print(f"🔁 [heartbeat] 连接中断，正在重连...")
            await asyncio.sleep(0.1)
            try:
                reader, writer = await asyncio.wait_for(
                    asyncio.open_connection(SERVER_IP, SERVER_PORT),
                    timeout=TIMEOUT
                )
                print(f"✅ [heartbeat] 重连成功")
                last_heartbeat = time.time()
            except Exception as e:
                print(f"❌ [heartbeat] 重连失败: {e}")
                return False, (time.time() - start_time) * 1000, total_bytes

        except asyncio.TimeoutError:
            print(f"⏰ [heartbeat] 网络超时，正在重连...")
            await asyncio.sleep(1)
            continue

        except Exception as e:
            print(f"❌ [heartbeat] 未知错误: {type(e).__name__}: {e}")
            return False, (time.time() - start_time) * 1000, total_bytes

    return True, (time.time() - start_time) * 1000, total_bytes


# 模式映射表
MODE_BEHAVIOR_MAP = {
	'request-response': behavior_request_response,
	'echo-loop': behavior_echo_loop,
	'streaming-receive': behavior_streaming_receive,
	'streaming-send': behavior_streaming_send,
	'connect-and-close': behavior_connect_and_close,
	'heartbeat-long-connection': behavior_heartbeat_long_connection,
}


async def tcp_request() -> Tuple[bool, float, int]:
	"""
    执行一次 TCP 请求，根据 TEST_MODE 调用不同行为
    Returns: (success, latency_ms, bytes_received)
    """
	try:
		reader, writer = await asyncio.wait_for(
			asyncio.open_connection(SERVER_IP, SERVER_PORT),
			timeout=TIMEOUT
		)
		
		behavior_func = MODE_BEHAVIOR_MAP.get(TEST_MODE)
		if not behavior_func:
			print(f"❌ 未知模式: {TEST_MODE}", file=sys.stderr)
			return False, 0, 0
		
		success, latency_ms, bytes_recv = await behavior_func(reader, writer)
		
		# 关闭连接
		try:
			writer.close()
			await writer.wait_closed()
		except:
			pass  # 忽略关闭异常
		
		return success, latency_ms, bytes_recv
	
	except asyncio.TimeoutError:
		print(f"⏰ 连接超时", file=sys.stderr)
		return False, 0, 0
	except ConnectionRefusedError:
		print(f"⛔ 连接被拒绝: 请检查 IP 和端口是否正确，或服务是否开启", file=sys.stderr)
		return False, 0, 0
	except ConnectionResetError:
		print(f"💥 连接被重置", file=sys.stderr)
		return False, 0, 0
	except Exception as e:
		print(f"❌ 未知错误 [{type(e).__name__}]: {e}", file=sys.stderr)
		return False, 0, 0


async def worker():
	global running, current_success, current_fail, current_latencies, current_bytes
	while running:
		success, latency_ms, bytes_recv = await tcp_request()
		async with lock:
			if success:
				current_success += 1
				current_latencies.append(latency_ms)
				current_bytes += bytes_recv
			else:
				current_fail += 1
		await asyncio.sleep(0.001)


async def print_stats():
	global current_success, current_fail, current_latencies, current_bytes
	global total_success, total_fail, total_latencies, total_bytes_received
	
	start_time = time.time()
	last_print = start_time
	
	print(f"🚀 压测开始 | 目标: {SERVER_IP}:{SERVER_PORT} | 模式: {TEST_MODE}")
	print(f"           | 并发: {NUM_WORKERS} 协程 | 每 {STAT_INTERVAL} 秒输出一次")
	print("-" * 80)
	print(f"{'时间':<8} {'QPS':<8} {'平均延迟(ms)':<12} {'P95延迟(ms)':<12} {'成功率':<8} {'接收数据'}")
	print("-" * 80)
	
	while time.time() - start_time < DURATION:
		await asyncio.sleep(STAT_INTERVAL - (time.time() - last_print) % STAT_INTERVAL)
		
		async with lock:
			duration = time.time() - last_print
			qps = current_success / duration if duration > 0 else 0
			avg_lat = statistics.mean(current_latencies) if current_latencies else 0
			p95_lat = statistics.quantiles(current_latencies, n=20)[-1] if len(current_latencies) >= 20 else 0
			success_rate = f"{current_success / (current_success + current_fail) * 100:.1f}%" if (
						                                                                                     current_success + current_fail) > 0 else "100%"
			
			print(f"{int(time.time() - start_time):>8}s "
			      f"{qps:>8.1f} "
			      f"{avg_lat:>12.2f} "
			      f"{p95_lat:>12.2f} "
			      f"{success_rate:>8} "
			      f"{current_bytes:,} B")
			
			total_success += current_success
			total_fail += current_fail
			total_latencies.extend(current_latencies)
			total_bytes_received += current_bytes
			
			current_success = 0
			current_fail = 0
			current_latencies.clear()
			current_bytes = 0
		
		last_print = time.time()


async def main():
	global running
	print("🚀 正在启动压测...", file=sys.stderr)
	print(f"🔧 模式: {TEST_MODE}", file=sys.stderr)
	
	if TEST_MODE not in MODE_BEHAVIOR_MAP:
		print(f"❌ 不支持的模式: {TEST_MODE}", file=sys.stderr)
		return
	
	stats_task = asyncio.create_task(print_stats())
	workers = [asyncio.create_task(worker()) for _ in range(NUM_WORKERS)]
	print(f"✅ 已启动 {len(workers)} 个 worker 协程", file=sys.stderr)
	
	await asyncio.sleep(DURATION)
	running = False
	print("🛑 正在停止所有 worker...", file=sys.stderr)
	
	await asyncio.gather(*workers, return_exceptions=True)
	await stats_task
	
	# =================== 最终报告 ===================
	total_time = DURATION
	total_requests = total_success + total_fail
	final_qps = total_success / total_time if total_time > 0 else 0
	final_avg_lat = statistics.mean(total_latencies) if total_latencies else 0
	final_p95_lat = statistics.quantiles(total_latencies, n=20)[-1] if len(total_latencies) >= 20 else 0
	final_p99_lat = statistics.quantiles(total_latencies, n=100)[-1] if len(total_latencies) >= 100 else 0
	throughput_mbps = (total_bytes_received * 8) / (total_time * 1024 * 1024) if total_time > 0 else 0
	
	print("=" * 80)
	print("📊 压测完成！最终性能报告")
	print("=" * 80)
	print(f"⏱  测试时长:        {total_time:.2f} 秒")
	print(f"🎯 总请求数:        {total_requests}")
	print(f"✅ 成功请求数:      {total_success}")
	print(f"❌ 失败请求数:      {total_fail}")
	print(f"🔥 平均 QPS:        {final_qps:.2f}")
	print(f"⚡ 平均响应时间:    {final_avg_lat:.2f} ms")
	print(f"📌 95% 响应时间:    {final_p95_lat:.2f} ms")
	print(f"📌 99% 响应时间:    {final_p99_lat:.2f} ms")
	print(f"📦 吞吐量:         {throughput_mbps:.2f} Mbps")
	print(f"📤 总接收字节数:    {total_bytes_received:,} 字节")
	print("=" * 80)


if __name__ == "__main__":
	try:
		asyncio.run(main())
	except KeyboardInterrupt:
		print("\n👋 压测被用户中断", file=sys.stderr)
		running = False
	except Exception as e:
		print(f"🚨 主程序异常: {e}", file=sys.stderr)